#include "ThreadWrapper.h"
#ifdef _WIN32
    #define WIN32_LEAN_AND_MEAN
    #include <Windows.h>
#else
    #include <sched.h>
    #include <pthread.h>
    #include <signal.h>
#endif

void ThreadWrapper::start() {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  if (m_isRunning.load(std::memory_order_acquire)) {
      std::cerr << "Thread is already running." << std::endl;
      return;
  }
  if (!m_mainFuncWithStop) {
      std::cerr << "No main function has been set." << std::endl;
      throw std::runtime_error("[ThreadWrapper]No main function has been set.");
  }

  m_stopRequested.store(false);

  std::thread tempThread([this]() {
      try {
          m_mainFuncWithStop(m_stopRequested);
      } catch (const std::exception& e) {
          std::cerr << "Exception in thread " << m_name << ": " << e.what() << std::endl;
      } catch (...) {
          std::cerr << "Unknown exception in thread " << m_name << std::endl;
      }
      m_isRunning.store(false);
      std::cout << "Thread " << m_name << " quit." << std::endl;
  });

  m_thread_ = std::move(tempThread);
  if (!m_thread_.joinable()) {
      std::cerr << "Failed to start thread." << std::endl;
      m_isRunning.store(false);
      m_isDetaching.store(false);
      throw std::runtime_error("[ThreadWrapper]Failed to start thread.");
  } 
  m_isRunning.store(true);
  try {
      __processTasks();
  }
  catch (const std::exception& e) {
      std::cerr << e.what() << '\n';
      return;
  }
  if (!m_name.empty()) {
      setName(m_name);        
  }
}
void ThreadWrapper::setSignalMask(const std::vector<int>& signals) {
  #ifndef _WIN32
    sigset_t signal_mask;
    sigemptyset(&signal_mask);
  
    for (const auto& sig : signals) {
        sigaddset(&signal_mask, sig);
    }
  
    int result = pthread_sigmask(SIG_BLOCK, &signal_mask, nullptr);
    if (result != 0) {
        throw std::runtime_error("[ThreadWrapper]Failed to set thread signal mask.");
    }
  #endif
}
void ThreadWrapper::stop(bool kill) {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  m_stopRequested.store(true, std::memory_order_release);
  m_stopRequested.notify_all();

  if (!kill) {
      if (m_thread_.joinable()) {
          m_thread_.join();
      } else {
          std::cerr << "Thread " << m_name << " was not joinable during stop()." << std::endl;
      }
  } else {
      std::cerr << "Warning: Forcibly killing thread " << m_name << ". This is dangerous!" << std::endl;
      #ifdef _WIN32
          TerminateThread(m_thread_.native_handle(), 0);
      #else
          pthread_cancel(m_thread_.native_handle());
      #endif
      m_isRunning.store(false);
  }

  m_stopRequested.store(false);
  if (kill || !m_thread_.joinable()) {
       if (m_isRunning.load()) {
           m_isRunning.store(false);
           std::cerr << "Thread " << m_name << ": Manually setting isRunning to false after kill/failed join." << std::endl;
       }
  }
}
void ThreadWrapper::detach() {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
    if (m_thread_.joinable()) {
        m_thread_.detach();
        m_isDetaching.store(true);
    }
}
void ThreadWrapper::setmain(std::function<void(const std::atomic<bool>&)> func) {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  m_mainFuncWithStop = func;
}
void ThreadWrapper::getlimit(size_t &stackSize) {
    this->__executeIfRunning([&]() {
        _getlimit(stackSize);
    });
}
void ThreadWrapper::setName(const std::string& name) {
    m_name = name;
    this->__executeIfRunning([&]() {
        _setName(name);
    });
}
void ThreadWrapper::setAffinity(const std::vector<int>& cpus) {
    this->__executeIfRunning([&]() {
        _setAffinity(cpus);
    });     
}
void ThreadWrapper::setPriority(int policy,int priority) {
    this->__executeIfRunning([&]() {
        _setPriority(policy,priority);
    });     
}
bool ThreadWrapper::isRunning() const {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  return this->m_isRunning.load();
}
bool ThreadWrapper::isDetaching() const {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  return this->m_isDetaching.load();
}
bool ThreadWrapper::hasMain() const {
  return m_mainFuncWithStop != nullptr;
}
bool ThreadWrapper::isJoinable() const {
  return m_thread_.joinable();
}
/*************** Private Functions ***************/
void ThreadWrapper::_setAffinity(const std::vector<int>& cpus) {
      unsigned int num_cores = std::thread::hardware_concurrency();
      if (num_cores == 0) {
          throw std::runtime_error("[ThreadWrapper]Failed to determine the number of available cores.");
      }
#ifdef _WIN32
      DWORD_PTR mask = 0;
      for (const auto& cpu : cpus) {
          if (cpu < 0 || static_cast<unsigned int>(cpu) >= num_cores) {
              throw std::runtime_error("[ThreadWrapper]Invalid core number specified.");
          }
          mask |= (1 << cpu);
      }

      DWORD_PTR result = SetThreadAffinityMask(m_thread_.native_handle(), mask);
      if (result == 0) {
          throw std::runtime_error("[ThreadWrapper]Failed to set thread CPU affinity.");
      }
#else
      cpu_set_t cpuset;
      CPU_ZERO(&cpuset);

      for (const auto& cpu : cpus) {
          if (cpu < 0 || static_cast<unsigned int>(cpu) >= num_cores) {
              throw std::runtime_error("Invalid core number specified.");
          }
          CPU_SET(cpu, &cpuset);
      }

      int result = pthread_setaffinity_np(m_thread_.native_handle(), sizeof(cpu_set_t), &cpuset);
      if (result != 0) {
          throw std::runtime_error("[ThreadWrapper]Failed to set thread CPU affinity.");
      }
#endif
}
void ThreadWrapper::_setPriority(int priority,int policy) {
#ifdef _WIN32
      if (priority < 1 || priority > 7) {
          throw std::runtime_error("[ThreadWrapper]Invalid priority value.");
      }
  
      int win_priority;
      switch (priority) {
          case 1: win_priority = THREAD_PRIORITY_IDLE; break;
          case 2: win_priority = THREAD_PRIORITY_LOWEST; break;
          case 3: win_priority = THREAD_PRIORITY_BELOW_NORMAL; break;
          case 4: win_priority = THREAD_PRIORITY_NORMAL; break;
          case 5: win_priority = THREAD_PRIORITY_ABOVE_NORMAL; break;
          case 6: win_priority = THREAD_PRIORITY_HIGHEST; break;
          case 7: win_priority = THREAD_PRIORITY_TIME_CRITICAL; break;
          default: throw std::runtime_error("[ThreadWrapper]Invalid priority value.");
      }
      if (!SetThreadPriority(reinterpret_cast<HANDLE>(m_thread_.native_handle()), win_priority)) {
          throw std::runtime_error("[ThreadWrapper]Failed to set thread priority.");
      }
#else
      if (policy != SCHED_FIFO && policy != SCHED_RR && policy != SCHED_OTHER) {
          throw std::runtime_error("[ThreadWrapper]Invalid scheduling policy.");
      }
  
      int min_priority = sched_get_priority_min(policy);
      int max_priority = sched_get_priority_max(policy);
      if (priority < min_priority || priority > max_priority) {
          throw std::runtime_error("[ThreadWrapper]Invalid priority value.");
      }
  
      sched_param sch_params;
      sch_params.sched_priority = priority;
  
      int result = pthread_setschedparam(m_thread_.native_handle(), policy, &sch_params);
      if (result != 0) {
          throw std::runtime_error("[ThreadWrapper]Failed to set thread priority.");
      }
      std::cout << "Thread priority: " << priority << std::endl;
#endif
}
void ThreadWrapper::_getlimit(size_t &stackSize) {
#ifdef _WIN32
    ULONG_PTR lowLimit, highLimit;
    GetCurrentThreadStackLimits(&lowLimit, &highLimit);
    int __stackSize = highLimit - lowLimit;
    std::cout << "Thread stack size: " << __stackSize << std::endl;
#else
    pthread_attr_t attr;
    if (pthread_getattr_np(m_thread_.native_handle(), &attr) == 0) {
        if (pthread_attr_getstacksize(&attr, &stackSize) == 0) {
            std::cout << "Thread stack size: " << stackSize << std::endl;
        } else {
            std::cerr << "Failed to get thread stack size." << std::endl;
        }
        pthread_attr_destroy(&attr);
    } else {
        std::cerr << "Failed to get thread attributes." << std::endl;
    }
#endif
} 
void ThreadWrapper::_setName(const std::string& name) {
#ifdef _WIN32
  std::wstring wname(m_name.begin(), m_name.end());
  SetThreadDescription(m_thread_.native_handle(), wname.c_str());
#else
  std::string short_name = m_name.substr(0, 15);
  int result = pthread_setname_np(m_thread_.native_handle(), short_name.c_str());
  if (result != 0) {
      throw std::runtime_error("[ThreadWrapper]Failed to set thread name.");
  }
#endif
}
inline void ThreadWrapper::__processTasks() {
    std::function<void()> task;
    while (true) {
        {
            std::lock_guard<std::mutex> lock(m_queueMutex);
            if (m_taskQueue.empty()) break;
            task = m_taskQueue.front();
            m_taskQueue.pop();
        }
        if (m_thread_.joinable()) {
            task();
        } else {
            std::lock_guard<std::mutex> lock(m_queueMutex);
            m_taskQueue.push(task);
            break;
        }
    }
}
inline void ThreadWrapper::__executeIfRunning(std::function<void()> func) {
  std::lock_guard<std::mutex> lock(m_startStopMutex);
  if (m_isRunning.load(std::memory_order_acquire)) {
      func();
  } else {
    std::lock_guard<std::mutex> lock(m_queueMutex);
    m_taskQueue.push(func);
    std::cerr << "Thread is not running. Task queued." << std::endl;
  }
}


